package com.liang.ali.yolox;

import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.index.NDIndex;
import ai.djl.ndarray.types.DataType;
import ai.djl.ndarray.types.Shape;
import ai.djl.nn.Activation;
import com.liang.ali.yolox.entity.FloatObject;
import io.grpc.ManagedChannel;
import io.grpc.ManagedChannelBuilder;
import net.coobird.thumbnailator.Thumbnails;
import org.tensorflow.framework.TensorProto;
import org.tensorflow.framework.TensorShapeProto;
import tensorflow.serving.Model;
import tensorflow.serving.Predict;
import tensorflow.serving.PredictionServiceGrpc;

import java.awt.image.BufferedImage;
import java.awt.image.Raster;
import java.nio.file.Paths;
import java.util.*;

import static com.liang.ali.yolox.PainImage.painImage;
import static org.tensorflow.framework.DataType.DT_FLOAT;


/*
 * @Author iLy
 * @Description
 * @ClassName test
 * @Date 2021/11/1 14:43
 * @VERSION 1.0
 **/
public class YoloxTest {

    public List<FloatObject> receiveCap(String fileName, String outputFolder) throws Exception {
        //String canvasPhotoBase64Data = "";

        //String file = base64ToFile(canvasPhotoBase64Data, "test.jpg").toString();
        String file = fileName;
        //读取文件，强制修改图片大小，设置输出文件格式bmp(模型定义时输入数据是无编码的)
        BufferedImage im = Thumbnails.of(file).forceSize(640, 640).outputFormat("bmp").asBufferedImage();
        //转换图片到图片数组，匹配输入数据类型为Float
        Raster raster = im.getData();
        List<Float> floatList = new ArrayList<>();
        float[] temp = new float[raster.getWidth() * raster.getHeight() * raster.getNumBands()];
        float[] pixels = raster.getPixels(0, 0, raster.getWidth(), raster.getHeight(), temp);
        for (float pixel : pixels) {
            floatList.add(pixel/255);
        }
        for (int i = 0; i < floatList.size(); i++) {
            if (i % 3 == 0){
                Float aFloat = floatList.get(i);
                aFloat -= 0.485f;
                aFloat /= 0.229f;
                floatList.set(i, aFloat);
            }else if (i % 3 == 1){
                Float aFloat = floatList.get(i);
                aFloat -= 0.456f;
                aFloat /= 0.224f;
                floatList.set(i, aFloat);
            }else {
                Float aFloat = floatList.get(i);
                aFloat -= 0.406f;
                aFloat /= 0.225f;
                floatList.set(i, aFloat);
            }
        }

        //#记个时
        long t = System.currentTimeMillis();
        //创建连接，注意usePlaintext设置为true表示用非SSL连接
        ManagedChannel channel = ManagedChannelBuilder.forAddress("121.196.146.141", 8502).usePlaintext(true).build();
        //这里还是先用block模式
        PredictionServiceGrpc.PredictionServiceBlockingStub stub = PredictionServiceGrpc.newBlockingStub(channel);
        //创建请求
        Predict.PredictRequest.Builder predictRequestBuilder = Predict.PredictRequest.newBuilder();
        //模型名称和模型方法名预设
        Model.ModelSpec.Builder modelSpecBuilder = Model.ModelSpec.newBuilder();
        modelSpecBuilder.setName("simple_test1");
        modelSpecBuilder.setSignatureName("serving_default");
        predictRequestBuilder.setModelSpec(modelSpecBuilder);
        //设置入参,访问默认是最新版本，如果需要特定版本可以使用tensorProtoBuilder.setVersionNumber方法
        TensorProto.Builder tensorProtoBuilder = TensorProto.newBuilder();
        tensorProtoBuilder.setDtype(DT_FLOAT);
        TensorShapeProto.Builder tensorShapeBuilder = TensorShapeProto.newBuilder();
        tensorShapeBuilder.addDim(TensorShapeProto.Dim.newBuilder().setSize(1));
        //#150528 = 224 * 224 * 3
        tensorShapeBuilder.addDim(TensorShapeProto.Dim.newBuilder().setSize(640));
        tensorShapeBuilder.addDim(TensorShapeProto.Dim.newBuilder().setSize(640));
        tensorShapeBuilder.addDim(TensorShapeProto.Dim.newBuilder().setSize(3));

        tensorProtoBuilder.setTensorShape(tensorShapeBuilder.build());
        tensorProtoBuilder.addAllFloatVal(floatList);
        predictRequestBuilder.putInputs("input_1", tensorProtoBuilder.build());
        //访问并获取结果
        Predict.PredictResponse predictResponse = stub.predict(predictRequestBuilder.build());


        Map<String, TensorProto> outputsMap = predictResponse.getOutputsMap();
        TensorProto concatenate_13 = outputsMap.get("concatenate_13");
        TensorProto concatenate_14 = outputsMap.get("concatenate_14");
        TensorProto concatenate_15 = outputsMap.get("concatenate_15");

        NDArray feat80ND = extracted(concatenate_13);
        NDArray feat40ND = extracted(concatenate_14);
        NDArray feat20ND = extracted(concatenate_15);

        float[][][][] feat80List = extracted1(concatenate_13);
        float[][][][] feat40List = extracted1(concatenate_14);
        float[][][][] feat20List = extracted1(concatenate_15);

        List<FloatObject> decode = decode(feat80List, feat40List, feat20List, feat80ND, feat40ND, feat20ND, file, outputFolder);

        return decode;
    }

    public List<FloatObject> decode(float[][][][] conv2d_109,float[][][][] conv2d_101,float[][][][] conv2d_93 ,NDArray conv2d109ND, NDArray conv2d101ND, NDArray conv2d93ND, String file, String outputFolder) throws Exception {
        List<String> classes = Arrays.asList("can", "bottle", "branch", "hay", "small-white", "grass",
                "leaf", "plastic-bag", "cup", "ship", "person", "the-nets", "water-float", "flower",
                "lotus-leaf", "paper", "black-floaters", "milk-box", "bubble", "bigunknow", "float1",
                "float2", "distance", "ball");
        int classesNum = classes.size();

        long[] inputShape = new long[]{640, 640};

        float iou = 0.3f;

        Image imgArr = ImageFactory.getInstance().fromFile(Paths.get(file));
        int height = imgArr.getHeight();
        int width = imgArr.getWidth();

        double[] model_image_size = new double[]{height,width};

        List<NDArray> feats = new ArrayList<>();
        feats.add(conv2d109ND);
        feats.add(conv2d101ND);
        feats.add(conv2d93ND);

        List<long[]> hw = new ArrayList<>();
        for (NDArray feat : feats) {
            Shape shape = feat.getShape();
            long l = shape.get(1);
            long l1 = shape.get(2);
            //grid_shape
            long[] gridShape = new long[]{l,l1};
            hw.add(gridShape);
        }

        NDManager ndManagerFeat = NDManager.newBaseManager();

        NDArray output = feats.get(0);
        output = output.reshape(new Shape(1, -1, 5 + classesNum));
        for (int i = 1; i < feats.size(); i++) {
            NDArray reshape = feats.get(i).reshape(new Shape(1, -1, 5 + classesNum));
            output = output.concat(reshape, 1);
        }
        float[] floats3 = output.toFloatArray();


        //grids shape strides
        List<NDArray> gridList = new ArrayList<>();
        List<long[]> shapeList = new ArrayList<>();
        List<NDArray> strideList = new ArrayList<>();

        for (int i = 0; i < hw.size(); i++) {
            //grid_y
            NDArray arangeY = ndManagerFeat.arange(hw.get(i)[1]);
            NDArray reshapeY = arangeY.reshape(new Shape(-1, 1));
            long[] repeatsY = new long[]{1, hw.get(i)[1]};
            NDArray gridY = reshapeY.tile(repeatsY);

            float[] gridY1 = gridY.toFloatArray();
            System.out.println(gridY1);

            //grid_x
            NDArray arangeX = ndManagerFeat.arange(hw.get(i)[0]);
            NDArray reshapeX = arangeX.reshape(new Shape(1, -1));
            long[] repeatsX = new long[]{hw.get(i)[1], 1};
            NDArray gridX = reshapeX.tile(repeatsX);

            float[] gridX1 = gridX.toFloatArray();
            System.out.println(gridX1);

            //grid
            NDArray grid = gridX.stack(gridY, 2).reshape(new Shape(1, -1, 2));
            float[] floats = grid.toFloatArray();
            System.out.println(floats);

            gridList.add(grid);

            //shape
            Shape gridShape = grid.getShape();
            long l = gridShape.get(0);
            long l1 = gridShape.get(1);

            NDArray ones = ndManagerFeat.ones(new Shape(l, l1, 1));
            NDArray mul = ones.mul(inputShape[0]);
            NDArray stride = mul.div(hw.get(i)[0]);
            float[] floats1 = stride.toFloatArray();
            strideList.add(stride);
        }

        //grids strides
        NDArray grids = gridList.get(0);
        NDArray strides = strideList.get(0);
        for (int i = 1; i < gridList.size(); i++) {
            grids = grids.concat(gridList.get(i), 1);
            strides = strides.concat(strideList.get(i), 1);
        }

        //box_xy box_wh
        NDArray add = output.get(new NDIndex(":, :, :2")).add(grids);
        NDArray mul = add.mul(strides);
        NDArray inputShapeN = ndManagerFeat.create(inputShape);
        NDArray boxXy = mul.div(inputShapeN.get(new NDIndex("-1")));
        float[] boxxyfloats = boxXy.toFloatArray();
        System.out.println(boxxyfloats);

        NDArray exp = output.get(new NDIndex(":, :, 2:4")).exp();
        NDArray mul1 = exp.mul(strides);
        NDArray boxWh = mul1.div(inputShapeN.get(new NDIndex("-1")));
        float[] boxWhfloats = boxWh.toFloatArray();
        System.out.println(boxWhfloats);

        //box_confidence box_class_probs
        NDArray boxConfidence = Activation.sigmoid(output.get(new NDIndex(":, :, 4:5")));
        NDArray boxClassProbs = Activation.sigmoid(output.get(new NDIndex(":, :, 5:")));
        float[] floats4 = boxConfidence.toFloatArray();
        float[] floats5 = boxClassProbs.toFloatArray();

        Map<String, NDArray> result = new HashMap<>();
        result.put("boxXy", boxXy);
        result.put("boxWh", boxWh);
        result.put("boxConfidence", boxConfidence);
        result.put("boxClassProbs", boxClassProbs);

        //imageShape
        NDArray imageShape = ndManagerFeat.create(model_image_size);

        NDArray boxes = yoloCorrectBoxes(result, inputShapeN, imageShape);

        NDArray boxScores = boxConfidence.mul(boxClassProbs);

        //mask
        NDArray mask = boxScores.gt(0.5f);
        boolean[] floats1 = mask.toBooleanArray();

        List<Float> boxScoreList = new ArrayList<>();
        List<float[]> boxList = new ArrayList<>();
        List<Integer> classList = new ArrayList<>();

        List<Integer> classResult = new ArrayList<>();
        List<float[]> boxResult = new ArrayList<>();
        List<Float> scoreResult = new ArrayList<>();

        long size1 = mask.size(1);
        for (long i = 0; i < size1; i++) {
            boolean[] booleans1 = mask.get(0).get(i).toBooleanArray();
            for (int l = 0; l < booleans1.length; l++) {
                if (booleans1[l]){
                    classList.add(l);
                    NDArray ndArray4 = boxes.get(0).get(i);
                    NDArray ndArray2 = ndArray4.toType(DataType.FLOAT32,true);
                    float[] floats = ndArray2.toFloatArray();
                    boxList.add(floats);
                    for (float aFloat : floats) {
                        System.out.println(aFloat);
                    }
                    //score
                    NDArray ndArray3 = boxScores.get(0).get(i);
                    NDArray ndArray5 = ndArray3.toType(DataType.FLOAT32, true);
                    float[] floats2 = ndArray5.toFloatArray();
                    for (float aFloat : floats2) {
                        if (aFloat < 0.5){
                            continue;
                        }else {
                            boxScoreList.add(aFloat);
                            System.out.println(aFloat);
                        }
                    }
                }
            }
        }

        for (int i = 0; i < classesNum; i++) {
            List<float[]> boxMap = new ArrayList<>();
            List<Float> scoreMap = new ArrayList<>();
            for (int i1 = 0; i1 < classList.size(); i1++) {
                if (classList.get(i1) == i){
                    boxMap.add(boxList.get(i1));
                    scoreMap.add(boxScoreList.get(i1));
                }
            }
            if (boxMap.size() > 0){
                List<Integer> list = single_class_non_max_suppression(boxMap, scoreMap, 0.5f, boxMap.size(), iou);
                for (Integer in : list) {
                    classResult.add(i);
                    boxResult.add(boxMap.get(in));
                    scoreResult.add(scoreMap.get(in));
                }
            }
        }

        Integer[] classArray = classResult.toArray(new Integer[classResult.size()]);
        float[][] boxArray = boxResult.toArray(new float[boxResult.size()][]);
        Float[] scoreArray = scoreResult.toArray(new Float[scoreResult.size()]);

        painImage(file, classArray, scoreArray, boxArray, outputFolder);

        List<FloatObject> objects = new ArrayList<>();
        for (int i = 0; i < classResult.size(); i++) {
            FloatObject floatObject = new FloatObject(scoreResult.get(i), classResult.get(i), boxResult.get(i));
            objects.add(floatObject);
        }
        return objects;
    }

    public class Idxs{
        int index;
        float probability;

        public int getIndex() {
            return index;
        }

        public void setIndex(int index) {
            this.index = index;
        }

        public float getProbability() {
            return probability;
        }

        public void setProbability(float probability) {
            this.probability = probability;
        }

        public Idxs(int index, float probability) {
            this.index = index;
            this.probability = probability;
        }
    }

    public List<Integer> single_class_non_max_suppression(List<float[]> bboxes, List<Float> confidences, float CONF_THRESH, int NUM_DETECTIONS, float IOU_THRESH) {
        if (bboxes.size() == 0)
            return null;
        List<Integer> conf_keep_idx = new ArrayList<>();//保存置信度大于CONF_THRESH的元素的下标
        List<Float> new_confidences = new ArrayList<>();//保存置信度大于CONF_THRESH的元素的值
        for (int i = 0; i < confidences.size(); i++) {
            float confidence = confidences.get(i);
            if (confidence > CONF_THRESH){
                conf_keep_idx.add(i);
                new_confidences.add(confidence);
            }
        }
        if (conf_keep_idx.isEmpty())
            return null;

        int alive_idx_size = conf_keep_idx.size();

        List<Idxs> idxsList = new ArrayList<>();//将置信度与下标对应
        for (int i = 0; i < alive_idx_size; i++) {
            Idxs idxs = new Idxs(conf_keep_idx.get(i), new_confidences.get(i));
            idxsList.add(idxs);
        }
        //按score升序排列
        Collections.sort(idxsList, new Comparator<Idxs>() {
            @Override
            public int compare(Idxs o1, Idxs o2) {
                return (int)(o1.getProbability()*10 - o2.getProbability()*10);
            }
        });

        float overlap_xmin, overlap_ymin, overlap_xmax, overlap_ymax;
        float overlap_w, overlap_h, overlap_area, overlap_ratio;
        //取出得分最高的bbox，计算剩下的bbox与它的交并比iou，去掉大于iou_thresh的bbox
        List<Integer> pickList = new ArrayList<>();
        while (idxsList.size() > 0) {
            int last = idxsList.size() - 1;
            if(pickList.size() >= NUM_DETECTIONS)//取置信度最高的NUM_DETECTIONS个结果
                break;
            int last_index = idxsList.get(last).getIndex();
            float last_area = (bboxes.get(last_index)[2] - bboxes.get(last_index)[0]) * (bboxes.get(last_index)[3] - bboxes.get(last_index)[1]);//area=(xmax-xmin)*(ymax-ymin)
            pickList.add(last_index);
            List<Idxs> idxs_to_remove = new ArrayList<>();//交并比过大需要移除的bbox
            for (int i = 0; i < last; i++) {
                int i_index = idxsList.get(i).getIndex();
                overlap_xmin = Math.max(bboxes.get(last_index)[0], bboxes.get(i_index)[0]);
                overlap_ymin = Math.max(bboxes.get(last_index)[1], bboxes.get(i_index)[1]);
                overlap_xmax = Math.min(bboxes.get(last_index)[2], bboxes.get(i_index)[2]);
                overlap_ymax = Math.min(bboxes.get(last_index)[3], bboxes.get(i_index)[3]);
                overlap_w = Math.max(0, overlap_xmax - overlap_xmin);
                overlap_h = Math.max(0, overlap_ymax - overlap_ymin);
                overlap_area = overlap_w * overlap_h;
                float i_area = (bboxes.get(i_index)[2] - bboxes.get(i_index)[0]) * (bboxes.get(i_index)[3] - bboxes.get(i_index)[1]);
                overlap_ratio = overlap_area / ( last_area + i_area - overlap_area);//IOU
                if (overlap_ratio > IOU_THRESH)
                    idxs_to_remove.add(idxsList.get(i));
            }
            idxs_to_remove.add(idxsList.get(last));
            idxsList.removeAll(idxs_to_remove);
        }
        return pickList;
    }


    public NDArray yoloCorrectBoxes( Map<String, NDArray> ptNDArraysMap, NDArray input_shape, NDArray image_shape){

        //boxXy  boxWh  boxConfidence  boxClassProbs
        NDArray boxXy = ptNDArraysMap.get("boxXy");
        NDArray boxWh = ptNDArraysMap.get("boxWh");
        NDArray boxConfidence = ptNDArraysMap.get("boxConfidence");
        NDArray boxClassProbs = ptNDArraysMap.get("boxClassProbs");

        //boxYx  boxHw
        NDArray boxYx = boxXy.flip(2);
        NDArray boxHw = boxWh.flip(2);
        float[] boxXy1 = boxXy.toFloatArray();
        float[] boxWh1 = boxWh.toFloatArray();
        float[] boxYx1 = boxYx.toFloatArray();
        float[] boxHw1 = boxHw.toFloatArray();


        //newShape
        NDArray min = input_shape.div(image_shape).min();
        NDArray mul = image_shape.mul(min);
        NDArray newShape = mul.round();
        NDArray ndArray6 = newShape.toType(DataType.FLOAT32, true);
        float[] newShape1 = ndArray6.toFloatArray();

        //offset
        NDArray sub = input_shape.sub(newShape);
        NDArray div = sub.div(2.);
        NDArray offset = div.div(input_shape);
        NDArray ndArray5 = offset.toType(DataType.FLOAT32, true);
        float[] offset1 = ndArray5.toFloatArray();

        //scale
        NDArray scale = input_shape.div(newShape);
        NDArray ndArray4 = scale.toType(DataType.FLOAT32, true);
        float[] scale1 = ndArray4.toFloatArray();

        //boxYxRe
        NDArray sub1 = boxYx.sub(offset);
        NDArray boxYxRe = sub1.mul(scale);
        NDArray ndArray3 = boxYxRe.toType(DataType.FLOAT32, true);
        float[] boxYxRe1 = ndArray3.toFloatArray();

        //boxHwRe
        NDArray boxHwRe = boxHw.mul(scale);
        NDArray ndArray = boxHwRe.toType(DataType.FLOAT32, true);
        float[] boxHwRe1 = ndArray.toFloatArray();

        //boxMins
        NDArray div1 = boxHwRe.div(2.0);
        NDArray boxMins = boxYxRe.sub(div1);
        NDArray ndArray2 = boxMins.toType(DataType.FLOAT32, true);
        float[] boxMin1 = ndArray2.toFloatArray();

        //boxMaxes
        NDArray div2 = boxHwRe.div(2.0);
        NDArray boxMaxes = boxYxRe.add(div2);
        NDArray ndArray1 = boxMaxes.toType(DataType.FLOAT32, true);
        float[] boxMaxes1 = ndArray1.toFloatArray();
        //boxes
        NDArray boxes = boxMins.concat(boxMaxes,2);
        NDArray ndArray8 = boxes.toType(DataType.FLOAT32, true);
        float[] boxes1 = ndArray8.toFloatArray();
        //boxesRe
        NDArray imageShapeCon = image_shape.concat(image_shape);
        NDArray boxesRe = boxes.mul(imageShapeCon);
        NDArray ndArray7 = boxesRe.toType(DataType.FLOAT32, true);
        float[] boxesRe1 = ndArray7.toFloatArray();

        return boxesRe;
    }


    private float[][][][] extracted1(TensorProto conv2d_1011) {
        List<Float> conv2d_101 = conv2d_1011.getFloatValList();
        TensorShapeProto conv2d_1011Shape = conv2d_1011.getTensorShape();
        List<TensorShapeProto.Dim> conv2d_1011DimList = conv2d_1011Shape.getDimList();

        float[] conv2dArray = new float[conv2d_101.size()];
        System.out.println(conv2dArray);
        for (int i = 0; i < conv2dArray.length; i++) {
            conv2dArray[i] = conv2d_101.get(i);
        }

        int count = 1;
//        for (TensorShapeProto.Dim dim : conv2d_1011DimList) {
//            count *= (int) dim.getSize();
//            System.out.println(dim.getSize());
//        }
        int dim1 = (int)conv2d_1011DimList.get(0).getSize();
        int dim2 = (int)conv2d_1011DimList.get(1).getSize();
        int dim3 = (int)conv2d_1011DimList.get(2).getSize();
        int dim4 = (int)conv2d_1011DimList.get(3).getSize();

        float[][][][] reshape = new float[dim1][dim2][dim3][dim4];
        int index = 0;
        for (int i = 0; i < dim1; i++) {
            for (int i1 = 0; i1 < dim2; i1++) {
                for (int i2 = 0; i2 < dim3; i2++) {
                    for (int i3 = 0; i3 < dim4; i3++) {
                        reshape[i][i1][i2][i3] = conv2dArray[index];
                        index++;
                    }
                }
            }
        }
        return reshape;
    }

    private NDArray extracted(TensorProto conv2d_1011) {
        List<Float> conv2d_101 = conv2d_1011.getFloatValList();
        TensorShapeProto conv2d_1011Shape = conv2d_1011.getTensorShape();
        List<TensorShapeProto.Dim> conv2d_1011DimList = conv2d_1011Shape.getDimList();

        float[] conv2dArray = new float[conv2d_101.size()];
        System.out.println(conv2dArray);
        for (int i = 0; i < conv2dArray.length; i++) {
            conv2dArray[i] = conv2d_101.get(i);
        }

        int count = 1;
//        for (TensorShapeProto.Dim dim : conv2d_1011DimList) {
//            count *= (int) dim.getSize();
//            System.out.println(dim.getSize());
//        }
        int dim1 = (int)conv2d_1011DimList.get(0).getSize();
        int dim2 = (int)conv2d_1011DimList.get(1).getSize();
        int dim3 = (int)conv2d_1011DimList.get(2).getSize();
        int dim4 = (int)conv2d_1011DimList.get(3).getSize();
        count = dim1*dim2*dim3*dim4;


        NDManager ndManager = NDManager.newBaseManager();
        NDArray reshape = ndManager.create(conv2dArray, new Shape(dim1, dim2, dim3, dim4));
        System.out.println(reshape);
//        NDArray ones = ndManager.ones(new Shape(count));
//        ones.set(conv2dArray);
//        System.out.println(ones);
//        NDArray reshape = ones.reshape(dim1, dim2, dim3, dim4);
        return reshape;
    }

}



